import numpy as np
import os
import struct

def gen_golden_data_hardshrink():
    # 生成包含正负数的输入数据（范围-2到2）
    input_x = np.random.uniform(-2, 2, [8, 2048]).astype(np.float32)
    
    # 定义Hardshrink参数 λ
    lambda_val = 0.5
    
    # 计算Hardshrink结果：f(x) = x if |x| > λ else 0
    golden = np.where(np.abs(input_x) > lambda_val, input_x, 0).astype(np.float32)
    
    # 创建输入输出目录（确保目录存在）
    os.makedirs("./input", exist_ok=True)
    os.makedirs("./output", exist_ok=True)
    
    # 生成与 HardshrinkCustomTilingData 结构体匹配的 tiling 数据
    # struct HardshrinkCustomTilingData {
    #     uint32_t totalLength;  // 8 * 2048 = 16384
    #     uint32_t tileNum;      // 8
    #     float lambda_val;      // 0.5
    # };
    
    total_length = 8 * 2048  # 16384
    tile_num = 8
    
    # 使用 struct.pack 确保正确的内存布局
    # 'I' = unsigned int (uint32_t), 'f' = float
    tiling_data = struct.pack('IIf', total_length, tile_num, lambda_val)
    
    # 保存 tiling 数据
    with open("./input/input_tiling.bin", "wb") as f:
        f.write(tiling_data)
    
    # 保存输入和输出数据
    input_x.tofile("./input/input_x.bin")
    golden.tofile("./output/golden.bin")
    
    print(f"生成Hardshrink数据完成！")
    print(f"λ = {lambda_val}")
    print(f"输入数据形状: {input_x.shape}")
    print(f"输出数据形状: {golden.shape}")
    print(f"输入范围: [{input_x.min():.3f}, {input_x.max():.3f}]")
    print(f"输出范围: [{golden.min():.3f}, {golden.max():.3f}]")
    
    # 验证一些样本
    print("\n样本验证:")
    for i in range(5):
        x_sample = input_x[0, i]
        y_sample = golden[0, i]
        status = "保留" if abs(x_sample) > lambda_val else "置零"
        print(f"x={x_sample:7.3f}, |x|={abs(x_sample):5.3f}, f(x)={y_sample:7.3f} [{status}]")

if __name__ == "__main__":
    gen_golden_data_hardshrink()